# -*- coding: utf-8 -*-
"""
Created on Tue Sep 20 01:00:02 2016

@author: zhiqiang
"""

path = r"D:\2016试题\B\B题附件"
fileName = "genotype.dat"
f = open(path+'\\'+fileName)
genData = f.readlines()
j = 1
result = {}
weiDian = genData[0]#表头
weiDian = weiDian.replace('\n','')
weiDian = weiDian.split(' ')
length = len(weiDian)
for gen_i in genData:
    if j==1:
        j+=1
        continue
    gen_i = gen_i.replace('\n','')
    gen_split = gen_i.split(' ')
    for i in range(length):
        result.setdefault(weiDian[i],[])
        result[weiDian[i]].append(gen_split[i])
f.close()

def getRatingnew(result):
    res = {}
    shangzimu = ['AA','TT','CC','GG']
    for k in result.keys():
        res.setdefault(k,{'jiankang':{},'bujiankang':{}})
        i = 1
        for vi in result[k]:
            if i<=500:
                if vi in shangzimu:
                    res[k]['jiankang'].setdefault(vi[0],0)
                    res[k]['jiankang'][vi[0]] += 2
                else:
                    res[k]['jiankang'].setdefault(vi[0],0)
                    res[k]['jiankang'][vi[0]] +=1
                    res[k]['jiankang'].setdefault(vi[1],0)
                    res[k]['jiankang'][vi[1]] +=1
            else:
                if vi in shangzimu:
                    res[k]['bujiankang'].setdefault(vi[0],0)
                    res[k]['bujiankang'][vi[0]] += 2
                else:
                    res[k]['bujiankang'].setdefault(vi[0],0)
                    res[k]['bujiankang'][vi[0]] +=1
                    res[k]['bujiankang'].setdefault(vi[1],0)
                    res[k]['bujiankang'][vi[1]] +=1
            i += 1
    return res

ress_new = getRatingnew(result)

kafang_param = {1: 7.8794}

def getKaFang_3_new(ress,kafang=kafang_param):
    #返回卡芳分布大于的基因位
    reslist = []
    for key in ress.keys():
        res_kafang = 0
        a=[]
        b=[]
        for k in ress[key]['bujiankang'].keys():
            a.append(ress[key]['bujiankang'][k])
            if k not in ress[key]['jiankang'].keys():
                b.append(0)
            else:
                b.append(ress[key]['jiankang'][k])
        zongshu = sum(a)+sum(b)
        for i in range(len(a)):
            res_kafang += (zongshu*((a[i]-sum(a)*(a[i]+b[i])/zongshu)**2))/(sum(a)*(a[i]+b[i]))+(zongshu*((b[i]-sum(b)*(a[i]+b[i])/zongshu)**2))/(sum(b)*(a[i]+b[i]))
        if res_kafang>kafang[len(a)-1]:
#            if key=='rs12133956':
#                print(res_kafang)
            reslist.append(key)
    return reslist
second_ans_for = getKaFang_3_new(ress_new,kafang_param)